#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA
#

__docformat__ = 'restructuredtext'
from scipy.sparse.linalg import spsolve as spsolve
from utils.ProjOperators import oned_inner_for_pide
from numpy import zeros


class IMEXTimeStepper(object):

    """
    Implicit Explicit time stepper as described in:

        @article{PIDpaper,
           author =  {Liming Feng, Vadim Linetsky},
           title =   {Pricing Options in Jump-Diffusion Models: An Extrapolation Approach.},
           journal = {Options Research},
           volume =  2,
           year =    2008,
           pages =   {304-325},
        }


    The algorithm presented here differs to the one in the paper on the
    way boundary conditions are treated. Boundary conditions are assumed
    to be zero in the paper, here we handle time dependent boundary conditions.
    """

    __slots__ = [
        'm_dolog',
        'm_ftime',
        'm_jbop',
        'm_ival',
        'm_jump',
        'm_tol',
        'm_rbc',
        'm_sbop',
        'm_stiff',
        'm_stime',
        'm_lbc',
        'm_load',
        'm_mass',
        'm_mesh',
        'm_method',
        'm_mbop',
        'm_numsteps',
        'm_numstates',
        'm_quads',
        'm_sfunc',
        'm_xl',
        'm_xr',
        ]

    def __init__(self):
        """
        Implicit - Explicit time stepper.

        Variables:
            m_dolog (bool) ~ log results.
            m_tol (float) ~ tolerance for adaptive steps.
            m_mass ~ mass matrix.
            m_stiff ~ full pde.
            m_jump ~ convolution.
            m_lbc ~ left boundary condition ( function (x, t) )
            m_jbop ~ pde compensation for implicit solution.
            m_method ~ Euler, midpoint, extrapolation.
            m_rbc ~ right boundary condition.
            m_mbop ~ right compensation for implicit solution
            m_load ~ rhs of the system of equations.
            m_ival ~ initial value
            m_numstepes ~ number of time steps.
            m_numstages ~ number of internal states in extrapolation.
            m_ftime ~ time to expiry.
        """

        (self.m_xl, self.m_xr) = (None, None)
        (self.m_rbc, self.m_lbc) = (None, None)
        (self.m_load, self.m_ival) = (None, None)
        (self.m_mbop, self.m_sbop, self.m_jbop) = (None, None, None)

        (self.m_mass, self.m_stiff, self.m_jump) = (None, None, None)

        (self.m_method) = (None)
        (self.m_dolog, self.m_tol) = (False, 1.e-16)

        (self.m_stime, self.m_ftime) = (0, None)
        (self.m_numsteps, self.m_numstates) = (None, None)
        (self.m_sfunc, self.m_quads, self.m_mesh) = (None, None, None)

    def __euler(self, start, stop, nsteps):
        """ Perform 'niter' time steps and return. """

        deltat = (stop - start) / nsteps

        fwdmat = self.m_mass.todense() + deltat * self.m_jump  # dense matrix
        implsmat = self.m_mass - deltat * self.m_stiff  # sparse matrix

        ival0 = self.m_ival.copy()
        lbc_prev = self.m_lbc(self.m_xl, start)
        rbc_prev = self.m_rbc(self.m_xr, start)

        (slbop, srbop) = (deltat * self.m_sbop[:,
                          0].reshape(ival0.shape), deltat
                          * self.m_sbop[:, 1].reshape(ival0.shape))

        (mlbop, mrbop) = (self.m_mbop[:, 0].reshape(ival0.shape),
                          self.m_mbop[:, 1].reshape(ival0.shape))

        (jlbop, jrbop) = (deltat * self.m_jbop[:,
                          0].reshape(ival0.shape), deltat
                          * self.m_jbop[:, 1].reshape(ival0.shape))

        for i in range(0, self.m_numsteps):

            # Fwd

            ival0 = fwdmat * ival0 + deltat * self.__inner_load(start,
                    i, deltat)

            # Boundary contribution.

            time_k = deltat * (i + 1) + start
            lbc_current = self.m_lbc(self.m_xl, time_k)
            rbc_current = self.m_rbc(self.m_xr, time_k)

            # mass

            ival0 += (lbc_prev - lbc_current) * mlbop + (rbc_prev
                    - rbc_current) * mrbop

            # jump

            ival0 += deltat * (lbc_prev * jlbop + rbc_prev * jrbop)

            # stiff

            ival0 += deltat * (lbc_current * slbop + rbc_current
                               * srbop)

            lbc_prev = lbc_current
            rbc_prev = rbc_current

            # Implicit

            ival0 = spsolve(implmat, ival0).reshape(self.m_ival.shape)
        return ival0

    def __inner_load(self, start, step, deltat):
        """ proyect load in finite element space at time t. """

        if self.m_load is None:
            cload = zeros(self.m_ival.shape)

        curr_time = start + step * deltat
        next_time = start + (step + 1) * deltat

        tmp = oned_inner_for_pide(
            self.m_sfunc,
            self.m_quads,
            self.m_mesh,
            self.m_load,
            curr_time,
            next_time,
            )

        return tmp[1:-1]


    def __midpoint(self, start, stop, nsteps):
        """ Implements the algorithm on page 318 of the paper. """

        deltat = (stop - start) / nsteps
        ival1 = self.__euler(start, deltat, 1)

        fwdmat = self.m_mass - deltat * self.m_stiff
        implmat = self.m_mass + deltat * self.m_stiff
        twojump = 2. * deltat * self.m_jump

        ival0 = self.m_ival.copy()

        lbc_prev = self.m_lbc(self.m_xl, deltat)
        rbc_prev = self.m_rbc(self.m_xr, deltat)

        (slbop, srbop) = (deltat * self.m_sbop[:,
                          0].reshape(ival0.shape), deltat
                          * self.m_sbop[:, 1].reshape(ival0.shape))

        (mlbop, mrbop) = (self.m_mbop[:, 0].reshape(ival0.shape),
                          self.m_mbop[:, 1].reshape(ival0.shape))

        (jlbop, jrbop) = (deltat * self.m_jbop[:,
                          0].reshape(ival0.shape), deltat
                          * self.m_jbop[:, 1].reshape(ival0.shape))

        for i in range(1, self.m_numsteps):

            # Fwd

            tmp = fwdmat * ival0 + twojump * ival1
            tmp += 2. * deltat * self.__inner_load(start,
                    i, deltat)

            # Boundary contribution.

            time_k = deltat * (i + 1) + start
            lbc_current = self.m_lbc(self.m_xl, time_k)
            rbc_current = self.m_rbc(self.m_xr, time_k)

            # mass

            tmp += (lbc_prev - lbc_current) * mlbop + (rbc_prev
                    - rbc_current) * mrbop

            ival1 = spsolve(implmat, tmp)

        return (ival0 + ival1) / 2.

    def get_startat(self):
        """ get starting time to simulate. """

        return self.m_stime


    def set_boperators(self, mbop, sbop, jbop):
        """
        This operators will be applied by the time stepper
        at each iteration. The time stepper will compute the
        solution at time 't' at the points 'xl' and 'xr',
        modify them and compensate the fwd step.
        """

        (self.m_mbop, self.m_sbop, self.m_jbop) = (mbop, sbop, jbop)

    def set_bcond(self, lbc, rbc):
        """
        Should be function of type (x, t).
        they will be evaluated as needed by the time stepper.
        """

        (self.m_lbc, self.m_rbc) = (lbc, rbc)

    def set_dolog(self, dolog):
        """
        log progress ???
        """

        self.m_dolog = dolog

    def set_ival(self, ival):
        """
        Initial value projected in the finite element space.
        """

        self.m_ival = ival

    def set_finite_elem_space(self, sfunc, quads, mesh):
        """ finite element space for time dependent load. """

        (self.m_sfunc, self.m_quads, self.m_mesh) = (sfunc, quads, mesh)

    def set_load(self, load):
        """
        Accounts for modifying the initial value. Load is assumed to be
        constant, time varying can be implemented too.
        """

        self.m_load = load

    def set_method(self, method):
        """ time stepping method. """

        assert method in ['euler', 'midpoint', 'extrapolation']
        self.m_method = method

    def set_moperators(self, mass, stiff, jump):
        """
        Three matrices are needed to perform the time steps:
        mass matrix, pde ~ stiffnes, jump ~ Toeplitz type.
        """

        (self.m_mass, self.m_stiff, self.m_jump) = (mass, stiff, jump)

    def set_number_steps(self, numsteps):
        """
        Will be used for non adaptive schemes.
        """

        self.m_numsteps = numsteps

    def set_numextrap_stages(self, numstages):
        """
        This number will define the max number of intermediate stages
        to go from step i to step i+1. Case the algorithm does not
        reach the desired tolerance, it will adjust the step size
        and repeat the process.
        """

        self.m_numstates = numstages

    def set_startat(self, stime):
        """ start simulating from this point on. """

        self.m_stime = stime

    def set_time2simulate(self, ftime):
        """ Termination time for simulation.  """

        self.m_ftime = ftime

    def set_xlxr(self, xleft, xright):
        """ Evaluate left and right boundary conditions here. """

        (self.m_xl, self.m_xr) = (xleft, xright)

    def step(self):
        """ Stiffness will be solve implicitly, non-stiffness is explicit. """

        assert 0 < self.m_ftime and 0 < self.m_numsteps
        assert self.m_sbop.shape == self.m_mbop.shape
        assert self.m_jbop.shape == self.m_mbop.shape

        assert self.m_mesh is not None
        assert self.m_sfunc is not None
        assert self.m_quads is not None

        # This iteration can be done more efficiently, the explicit
        # step can be sped up because M is diagonal and jump is Toeplitz.
        # Note that smat is sparse and fwdmat is diagonal + Toeplitz.

        if self.m_method is None:
            raise ValueError, 'set the time sttepping method.'

        if 'euler' == self.m_method:
            psolu = self.__euler(self.m_stime, self.m_ftime,
                                 self.m_numsteps)

        if 'midpoint' == self.m_method:
            psolu = self.__midpoint(self.m_stime, self.m_ftime,
                                    self.m_numsteps)

        solu = zeros((psolu.size + 2, 1), 'float64')
        solu[1:-1] = psolu
        solu[0] = self.m_lbc(self.m_xl, self.m_ftime)
        solu[-1] = self.m_rbc(self.m_xr, self.m_ftime)
        return solu


